XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python klein_gordon4d.py --model=pinn --equation=klein_gordon4d --nc=16 --nc_test=50 --seed=111 --lr=1e-3 --epochs=50000 --mlp=modified_mlp --n_layers=5 --features=128 --out_dim=1 --k=2 --log_iter=1000